"""
生成构件级仿真数据集
功能：模拟不同结构参数组合下的地震响应数据
输出：包含结构参数、地震波、损伤标签的CSV文件
"""
import numpy as np
import pandas as pd
from tqdm import tqdm  # 进度条工具
import json

class DataGenerator:
    def __init__(self, num_samples=10000):
        """
        参数说明：
        - num_samples: 生成样本总数（默认1万）
        """
        self.num_samples = num_samples
        # 结构参数范围（单位：mm）
        self.beam_width_range = (200, 600)  # 梁宽
        self.beam_height_range = (300, 800)  # 梁高
        self.concrete_grades = ['C30', 'C40', 'C50']  # 混凝土强度等级
        self.rebar_ratios = [0.01, 0.015, 0.02, 0.025, 0.03]  # 配筋率

    def _generate_waveform(self, duration=10, sample_rate=100):
        """
        生成模拟地震波（专业解释：非平稳随机振动信号）
        参数：
        - duration: 持续时间（秒）
        - sample_rate: 采样率（Hz）
        返回：形状为(3, 1000)的三轴加速度数据
        """
        t = np.linspace(0, duration, duration * sample_rate)
        # 主震成分（低频大振幅）
        main_freq = np.random.uniform(0.5, 5.0)  # 主频0.5-5Hz
        main_shock = 0.5 * np.sin(2 * np.pi * main_freq * t)
        # 余震成分（高频小振幅）
        after_shock = 0.2 * np.sin(2 * np.pi * 3 * main_freq * t)
        # 添加高斯噪声（模拟传感器噪声）
        noise = np.random.normal(0, 0.05, len(t))
        # 合成三轴数据（假设X轴为主方向）
        x = main_shock + after_shock + noise
        y = 0.7 * x + np.random.normal(0, 0.03, len(t))
        z = 0.3 * x + np.random.normal(0, 0.01, len(t))
        return np.stack([x, y, z], axis=0)  # 形状(3, 1000)

    def _generate_damage_label(self, params):
        """
        计算损伤标签（大白话：根据结构参数评估受损程度）
        参数：
        - params: 包含结构参数的字典
        返回：0~1之间的损伤指数
        """
        # 损伤影响因素：配筋率越低、梁高越大越易受损
        base_damage = (params['rebar_ratio'] - 0.01) / 0.02  # 配筋率影响
        height_effect = (params['beam_height'] - 300) / 500  # 梁高影响
        # 混凝土强度映射（C30→0.2, C40→0.1, C50→0.0）
        grade_map = {'C30': 0.2, 'C40': 0.1, 'C50': 0.0}
        damage = 0.5 * (1 - base_damage) + 0.3 * height_effect + 0.2 * grade_map[params['concrete_grade']]
        return np.clip(damage + np.random.normal(0, 0.05), 0, 1)  # 添加随机扰动

    def generate_dataset(self, output_path):
        """生成并保存数据集"""
        data = []
        for _ in tqdm(range(self.num_samples)):
            # 随机生成结构参数
            params = {
                'beam_width': np.random.randint(*self.beam_width_range),
                'beam_height': np.random.randint(*self.beam_height_range),
                'concrete_grade': np.random.choice(self.concrete_grades),
                'rebar_ratio': np.random.choice(self.rebar_ratios)
            }
            # 生成地震波和损伤标签
            wave = self._generate_waveform()
            damage = self._generate_damage_label(params)
            # 展平数据便于存储
            # record = {
            #     **params,
            #     'acc_x': wave[0].tolist(),
            #     'acc_y': wave[1].tolist(),
            #     'acc_z': wave[2].tolist(),
            #     'damage_index': damage
            # }
            # 修改后（正确方式）
            record = {
                # 使用JSON序列化确保数据格式
                'acc_x': json.dumps(wave[0].tolist()),
                'acc_y': json.dumps(wave[1].tolist()),
                'acc_z': json.dumps(wave[2].tolist()),
                # 其他参数保持不变...
                **params,
                'damage_index': damage
            }
            data.append(record)

        # 保存为CSV
        df = pd.DataFrame(data)
        df.to_csv(output_path, index=False)
        print(f"数据集已保存至 {output_path}，共 {len(df)} 条样本")


if __name__ == "__main__":
    generator = DataGenerator(num_samples=10000)
    generator.generate_dataset("data/component_data.csv")